#!/bin/bash

set -u

declare -A CMD_SRC_FUNC_ARRAY
CMD_SRC_FUNC_ARRAY["build"]=cmd_src_build
CMD_SRC_FUNC_ARRAY["install"]=cmd_src_install
CMD_SRC_FUNC_ARRAY["uninstall"]=cmd_src_uninstall
CMD_SRC_FUNC_ARRAY["reinstall"]=cmd_src_reinstall

L4T_TOOLCHAIN_TOP="${HOME}/nvidia/toolchain"
L4T_SRC_TOP=${HOME}/nvidia/source_code/l4t_${L4T_VER}/Linux_for_Tegra/kernel_src
L4T_SRC_BRANCH=master
L4T_SRC_OUTPUT_TOP=${HOME}/nvidia/source_code/l4t_${L4T_VER}/OUTPUT_${L4T_SRC_BRANCH}

declare -A L4T_TOOLCHAIN_ARRAY
_L4T_TOOLCHAIN_0="/usr/bin/aarch64-linux-gnu-"
_L4T_TOOLCHAIN_1="${L4T_TOOLCHAIN_TOP}/Toolchain_Linaro_GCC_7.3.1/gcc-linaro-7.3.1-2018.05-x86_64_aarch64-linux-gnu/bin/aarch64-linux-gnu-"
_L4T_TOOLCHAIN_2="${L4T_TOOLCHAIN_TOP}/Toolchain_Bootlin_GCC_9.3/aarch64--glibc--stable-final/bin/aarch64-linux-"
L4T_TOOLCHAIN_ARRAY["32"]=${_L4T_TOOLCHAIN_1}
L4T_TOOLCHAIN_ARRAY["35"]=${_L4T_TOOLCHAIN_2}
L4T_TOOLCHAIN=${L4T_TOOLCHAIN_ARRAY[${L4T_VER_MAJOR}]}

declare -A L4T_SRC_MAKE_TOP_ARRAY
_L4T_SRC_MAKE_TOP_1="${L4T_SRC_TOP}/kernel/kernel-4.9"
_L4T_SRC_MAKE_TOP_2="${L4T_SRC_TOP}/kernel/kernel-5.10"
L4T_SRC_MAKE_TOP_ARRAY["32"]=${_L4T_SRC_MAKE_TOP_1}
L4T_SRC_MAKE_TOP_ARRAY["35"]=${_L4T_SRC_MAKE_TOP_2}
L4T_SRC_MAKE_TOP=${L4T_SRC_MAKE_TOP_ARRAY[${L4T_VER_MAJOR}]}

function cmd_src()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	if [[ $# -le 0 ]]; then
		echo -e "[ERROR] Must specify a SRC SUB_CMD.\n"
		exit 1
	fi

	local SUB_CMD=$1
	shift

	local key
	local found=0
	for key in ${!CMD_SRC_FUNC_ARRAY[@]}; do
		if [[ ${SUB_CMD} == ${key} ]]; then
			found=1
		fi
	done
	if [[ ${found} -ne 1 ]]; then
		echo -e "[ERROR] Unsupport SRC SUB_CMD: ${SUB_CMD}\n"
		exit 1
	fi

	l4t_src_init

	${CMD_SRC_FUNC_ARRAY[${SUB_CMD}]} $@
}

function cmd_src_build()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local L4T_SRC_OUTPUT_VERMAGIC_FILE="${L4T_SRC_OUTPUT_TOP}/include/config/kernel.release"
	if [[ -f ${L4T_SRC_OUTPUT_VERMAGIC_FILE} ]]; then
		export LOCALVERSION="-$(cat ${L4T_SRC_OUTPUT_VERMAGIC_FILE} | cut -d - -f 2-)"
	else
		local RK_VER_SUFFIX=$(date +"%Y%m%d%H%M%S")
		export LOCALVERSION="-tegra-${RK_VER_SUFFIX}"
	fi
	echo -e "[DEBUG] LOCALVERSION=${LOCALVERSION}"

	local MAKE_CONFIG=tegra_ritrontek_defconfig
	echo -e "[DEBUG] MAKE_CONFIG=${MAKE_CONFIG}"
	if [[ ! -f "${L4T_SRC_OUTPUT_TOP}/.config" ]]; then
		make -C ${L4T_SRC_MAKE_TOP} \
			${MAKE_CONFIG} \
			O=${L4T_SRC_OUTPUT_TOP} \
			ARCH=arm64 \
			CROSS_COMPILE=${L4T_TOOLCHAIN}
	fi

	if [[ $# -le 0 ]]; then
		local target="all"
	else
		local target=$1
	fi

	case ${target} in
		"all" )
			l4t_src_build dtbs
			l4t_src_build Image
			l4t_src_build modules
			l4t_src_module_display build
			l4t_src_module_ngbe build
			l4t_src_module_ch36x build
			;;
		"dtbs" | "Image" | "modules" )
			l4t_src_build ${target}
			;;
		"display" )
			l4t_src_module_display build
			;;
		"ngbe" )
			l4t_src_module_ngbe build
			;;
		"ch36x" )
			l4t_src_module_ch36x build
			;;
		*)
			echo -e "[ERROR] Unsupport SRC BUILD CMD: ${SUB_CMD}\n"
			usage
			exit 1
			;;
	esac
}

function cmd_src_install()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local KERNEL_SRC_OUTPUT=${L4T_SRC_OUTPUT_TOP}/arch/arm64/boot/Image
	local MOD_SRC_OUTPUT=${L4T_SRC_OUTPUT_TOP}/drivers/net/can/mttcan/native/mttcan.ko
	if [[ ${L4T_VER_MAJOR} -lt 35 ]]; then
		local DTB_SRC_OUTPUT_TOP=${L4T_SRC_OUTPUT_TOP}/arch/arm64/boot/dts/
	else
		local DTB_SRC_OUTPUT_TOP=${L4T_SRC_OUTPUT_TOP}/arch/arm64/boot/dts/nvidia/
	fi

	if [[ ! -d ${L4T_SRC_OUTPUT_TOP} ]] \
		|| [[ ! -f ${KERNEL_SRC_OUTPUT} ]] \
		|| [[ ! -f ${MOD_SRC_OUTPUT} ]] \
		|| [[ ! -d ${DTB_SRC_OUTPUT_TOP} ]]; then
		echo -e "[ERROR] L4T SRC not fully build, please build first."
		exit 1
	fi

	local KERNEL=${L4T_TOP}/kernel/Image
	local KERNEL_SRC=${L4T_TOP}/kernel_src/Image
	if [[ $(readlink ${KERNEL}) == ${KERNEL_SRC} ]]; then
		echo -e "[INFO] L4T SRC already installed, no need to install again."
		return
	fi

	if [[ -d ${L4T_TOP}/kernel_src ]]; then
		rm -rf ${L4T_TOP}/kernel_src
	fi

	local dtb_src_top=${L4T_TOP}/kernel_src/dtb
	mkdir -p ${dtb_src_top}
	ln ${KERNEL_SRC_OUTPUT} ${KERNEL_SRC}
	ln -sf ${KERNEL_SRC} ${KERNEL}

	local dtb_src_output_list=$(ls ${DTB_SRC_OUTPUT_TOP}/{*.dtb,*.dtbo})
	local dtb_src_output
	for dtb_src_output in ${dtb_src_output_list}
	do
		ln ${dtb_src_output} ${dtb_src_top}
	done
	local dtb_src_list=$(ls ${dtb_src_top}/{*.dtb,*.dtbo})
	local dtb_src
	for dtb_src in ${dtb_src_list}
	do
		ln -sf ${dtb_src} ${L4T_TOP}/kernel/dtb/
	done

	l4t_src_modules_install
	l4t_src_module_display install
	l4t_src_module_ngbe install
	l4t_src_module_ch36x install

	sudo_cmd "depmod \
			-b ${L4T_TOP}/rootfs \
			-F ${L4T_SRC_OUTPUT_TOP}/System.map \
			$(get_cur_kernel_vermagic)"

	echo -e "\n[INFO] Install Done.\n"
}

function cmd_src_uninstall()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local KERNEL=${L4T_TOP}/kernel/Image
	local KERNEL_SRC=${L4T_TOP}/kernel_src/Image
	if [[ $(readlink ${KERNEL}) != ${KERNEL_SRC} ]]; then
		echo -e "[INFO] L4T SRC not installed, cannot uninstall."
		return
	fi

	l4t_src_module_ch36x uninstall
	l4t_src_module_ngbe uninstall
	l4t_src_module_display uninstall
	l4t_src_modules_uninstall

	local dtb_list=$(ls ${L4T_TOP}/kernel/dtb/{*.dtb,*.dtbo})
	local dtb
	for dtb in ${dtb_list}
	do
		if [[ $(dirname $(readlink ${dtb})) == ${L4T_TOP}/kernel_src/dtb ]]; then
			rm ${dtb}
		fi
	done
	local dtb_bsp_list=$(ls ${L4T_TOP}/kernel_bsp/dtb/{*.dtb,*.dtbo})
	local dtb_bsp
	for dtb_bsp in ${dtb_bsp_list}
	do
		ln -sf ${dtb_bsp} ${L4T_TOP}/kernel/dtb/
	done

	ln -sf ${L4T_TOP}/kernel_bsp/Image ${KERNEL}

	echo -e "\n[INFO] Uninstall Done.\n"
}

function cmd_src_reinstall()
{
	cmd_src_uninstall
	cmd_src_install
}

function l4t_src_init()
{
	if [[ -d ${L4T_SRC_TOP} ]]; then
		cd ${L4T_SRC_TOP}
		if [[ $(git branch --show-current) != ${L4T_SRC_BRANCH} ]]; then
			echo -e "[INFO] Switch Kernel SRC branch to: ${L4T_SRC_BRANCH}\n"
			git checkout ${L4T_SRC_BRANCH}
			local ret=$?
			if [[ ${ret} != 0 ]]; then
				echo -e "[ERROR] Switch Kernel SRC branch error.\n"
				exit 1
			fi
		fi
		cd - &> /dev/null
	else
		local L4T_SRC_GIT_URL="gitea@192.168.1.55:nvidia_jetson/l4t_${L4T_VER}_kernel_src.git"

		echo -e "[INFO] Kernel SRC directory doesnot exist: ${L4T_SRC_TOP}"
		echo -e "[INFO] Clone from Gitea: ${L4T_SRC_GIT_URL}\n"
		git clone -b ${L4T_SRC_BRANCH} ${L4T_SRC_GIT_URL} ${L4T_SRC_TOP}
		local ret=$?
		if [[ ${ret} != 0 ]]; then
			echo -e "[ERROR] Clone Kernel SRC error.\n"
			exit 1
		fi
	fi
}

function l4t_src_build()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local target=$1
	make -C ${L4T_SRC_MAKE_TOP} ${target} \
		-j$(nproc) \
		O=${L4T_SRC_OUTPUT_TOP} \
		ARCH=arm64 \
		CROSS_COMPILE=${L4T_TOOLCHAIN}

	local ret=$?
	if [[ ${ret} -ne 0 ]]; then
		echo -e "[ERROR] L4T build error!\n"
		exit 1
	fi
}

function l4t_src_modules_install()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local L4T_KERNEL_VERMAGIC=$(get_cur_kernel_vermagic)
	if [[ -d ${L4T_TOP}/rootfs/lib/modules/${L4T_KERNEL_VERMAGIC} ]]; then
		echo -e "[INFO] Kernel Modules already installed: ${L4T_KERNEL_VERMAGIC}"
	else
		echo -e "[INFO] Install Kernel Modules ..."
		sudo_cmd "make -C ${L4T_SRC_MAKE_TOP} \
			modules_install \
			O=${L4T_SRC_OUTPUT_TOP} \
			ARCH=arm64 \
			CROSS_COMPILE=${L4T_TOOLCHAIN} \
			INSTALL_MOD_PATH=${L4T_TOP}/rootfs"
	fi
}

function l4t_src_modules_uninstall()
{
	echo -e "[DEBUG] ${FUNCNAME} $@"

	local L4T_KERNEL_VERMAGIC=$(get_cur_kernel_vermagic)
	local RK_VER_SUFFIX=$(get_kernel_vermagic_rksuffix ${L4T_KERNEL_VERMAGIC})
	local L4T_SRC_MODULES_DIR=${L4T_TOP}/rootfs/lib/modules/${L4T_KERNEL_VERMAGIC}

	if [[ -n ${RK_VER_SUFFIX} ]] && [[ -d ${L4T_SRC_MODULES_DIR} ]]; then
		sudo_cmd "rm -rf ${L4T_SRC_MODULES_DIR}"
	fi
}

function l4t_src_module_display()
{
	echo -e "\n[DEBUG] ${FUNCNAME} $@"

	echo -e "[DEBUG] L4T_VER_MAJOR=${L4T_VER_MAJOR}"
	if [[ ${L4T_VER_MAJOR} -lt 35 ]]; then
		echo -e "[DEBUG] skip display module"
		return
	fi

	local L4T_MODULE_DISPLAY_SRC="${HOME}/nvidia/source_code/l4t_${L4T_VER}_nvidia_display_driver"
	local L4T_MODULE_DISPLAY="${HOME}/nvidia/source_code/l4t_${L4T_VER}_nvidia_display_driver/kernel-open/nvidia.ko"

	case $1 in
		"build")
			if [[ ! -d ${L4T_MODULE_DISPLAY_SRC} ]]; then
				git clone gitea@192.168.1.55:nvidia_jetson/l4t_${L4T_VER}_nvidia_display_driver.git ${L4T_MODULE_DISPLAY_SRC}
			fi

			if [[ ! -f ${L4T_MODULE_DISPLAY} ]] || \
				[[ -f ${L4T_MODULE_DISPLAY} && $(get_module_vermagic ${L4T_MODULE_DISPLAY}) != $(get_cur_kernel_vermagic) ]];
			then
				make clean \
					-C ${L4T_MODULE_DISPLAY_SRC} \
					-j$(nproc) \
					SYSSRC=${L4T_SRC_MAKE_TOP} \
					SYSOUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN} \
					TARGET_ARCH=aarch64 \
					CC=${L4T_TOOLCHAIN}gcc \
					LD=${L4T_TOOLCHAIN}ld.bfd \
					AR=${L4T_TOOLCHAIN}ar \
					CXX=${L4T_TOOLCHAIN}g++ \
					OBJCOPY=${L4T_TOOLCHAIN}objcopy

				make modules \
					-C ${L4T_MODULE_DISPLAY_SRC} \
					-j$(nproc) \
					SYSSRC=${L4T_SRC_MAKE_TOP} \
					SYSOUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN} \
					TARGET_ARCH=aarch64 \
					CC=${L4T_TOOLCHAIN}gcc \
					LD=${L4T_TOOLCHAIN}ld.bfd \
					AR=${L4T_TOOLCHAIN}ar \
					CXX=${L4T_TOOLCHAIN}g++ \
					OBJCOPY=${L4T_TOOLCHAIN}objcopy

				local BUILD_RET=$?
				if [[ ${BUILD_RET} -ne 0 ]]; then
					echo -e "[ERROR] L4T Module Display build error!\n"
					exit 1
				fi
			fi
			;;
		"install")
			if [[ ! -f ${L4T_MODULE_DISPLAY} ]]; then
				${FUNCNAME} build

				if [[ ! -f ${L4T_MODULE_DISPLAY} ]]; then
					echo -e "[ERROR] L4T Module Display install error!\n"
					echo -e "NVIDIA Display Module file does not exist: ${L4T_MODULE_DISPLAY}.\n"
					exit 1
				fi
			fi

			sudo_cmd "make modules_install \
				-C ${L4T_MODULE_DISPLAY_SRC} \
				-j$(nproc) \
				SYSSRC=${L4T_SRC_MAKE_TOP} \
				SYSOUT=${L4T_SRC_OUTPUT_TOP} \
				INSTALL_MOD_PATH=${L4T_TOP}/rootfs \
				ARCH=arm64 \
				CROSS_COMPILE=${L4T_TOOLCHAIN} \
				TARGET_ARCH=aarch64 \
				CC=${L4T_TOOLCHAIN}gcc \
				LD=${L4T_TOOLCHAIN}ld.bfd \
				AR=${L4T_TOOLCHAIN}ar \
				CXX=${L4T_TOOLCHAIN}g++ \
				OBJCOPY=${L4T_TOOLCHAIN}objcopy"
			;;
		"uninstall")
			;;
		*)
			echo -e "[ERROR] ${FUNCNAME} argument error!\n"
			exit 1
			;;
	esac
}

function l4t_src_module_ngbe()
{
	echo -e "\n[DEBUG] ${FUNCNAME} $@"

	local rk_mod_ngbe=$(conf1_get ${L4T_TOP}/${JETSON_CONF}.conf "RK_MOD_NGBE")
	echo -e "[DEBUG] RK_MOD_NGBE=${rk_mod_ngbe}"

	if [[ -z ${rk_mod_ngbe} ]] || [[ ${rk_mod_ngbe} != "Y" ]]; then
		echo -e "[DEBUG] skip ngbe module"
		return
	fi

	local NGBE_VER=1.2.3
	local L4T_KERNEL_VERMAGIC=$(get_cur_kernel_vermagic)
	local L4T_MODULE_NGBE_TOP="${HOME}/nvidia/source_code/ngbe-${NGBE_VER}"
	local L4T_MODULE_NGBE="${L4T_MODULE_NGBE_TOP}/src/ngbe.ko"
	local L4T_MODULE_NGBE_ROOTFS_DIR="${L4T_TOP}/rootfs/lib/modules/${L4T_KERNEL_VERMAGIC}/kernel/drivers/net/ethernet/ngbe"

	case $1 in
		"build")
			if [[ ! -d ${L4T_MODULE_NGBE_TOP} ]]; then
				git clone gitea@192.168.1.55:pan.kang/ngbe-${NGBE_VER}.git ${L4T_MODULE_NGBE_TOP}
			fi

			if [[ ! -f ${L4T_MODULE_NGBE} ]] || \
				[[ -f ${L4T_MODULE_NGBE} && $(get_module_vermagic ${L4T_MODULE_NGBE}) != ${L4T_KERNEL_VERMAGIC} ]];
			then
				make clean \
					-C ${L4T_MODULE_NGBE_TOP}/src \
					-j$(nproc) \
					KERNEL_SRC=${L4T_SRC_MAKE_TOP} \
					KERNEL_OUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN}

				make \
					-C ${L4T_MODULE_NGBE_TOP}/src \
					-j$(nproc) \
					KERNEL_SRC=${L4T_SRC_MAKE_TOP} \
					KERNEL_OUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN}

				local BUILD_RET=$?
				if [[ ${BUILD_RET} -ne 0 ]]; then
					echo -e "[ERROR] L4T Module NGBE build error!\n"
					exit 1
				fi
			fi
			;;
		"install")
			if [[ ! -f ${L4T_MODULE_NGBE} ]]; then
				${FUNCNAME} build

				if [[ ! -f ${L4T_MODULE_NGBE} ]]; then
					echo -e "[ERROR] L4T Module NGBE install error!\n"
					echo -e "NGBE Module file does not exist: ${L4T_MODULE_NGBE}.\n"
					exit 1
				fi
			fi

			echo -e "> ${L4T_MODULE_NGBE_ROOTFS_DIR}/$(basename ${L4T_MODULE_NGBE})"
			sudo_cmd "mkdir -p ${L4T_MODULE_NGBE_ROOTFS_DIR} \
				&& cp ${L4T_MODULE_NGBE} ${L4T_MODULE_NGBE_ROOTFS_DIR}"
			;;
		"uninstall")
			;;
		*)
			echo -e "[ERROR] ${FUNCNAME} argument error!\n"
			exit 1
			;;
	esac
}

function l4t_src_module_ch36x()
{
	echo -e "\n[DEBUG] ${FUNCNAME} $@"

	local rk_mod_ch36x=$(conf1_get ${L4T_TOP}/${JETSON_CONF}.conf "RK_MOD_CH36X")
	echo -e "[DEBUG] RK_MOD_CH36X=${rk_mod_ch36x}"

	if [[ -z ${rk_mod_ch36x} ]] || [[ ${rk_mod_ch36x} != "Y" ]]; then
		echo -e "[DEBUG] skip ch36x module"
		return
	fi

	local L4T_KERNEL_VERMAGIC=$(get_cur_kernel_vermagic)
	local L4T_MODULE_CH36X_TOP="${HOME}/nvidia/source_code/ch36x"
	local L4T_MODULE_CH36X="${L4T_MODULE_CH36X_TOP}/driver/ch36x.ko"
	local L4T_MODULE_CH36X_ROOTFS_DIR="${L4T_TOP}/rootfs/lib/modules/${L4T_KERNEL_VERMAGIC}/kernel/drivers/misc/"

	case $1 in
		"build")
			if [[ ! -d ${L4T_MODULE_CH36X_TOP} ]]; then
				git clone gitea@192.168.1.55:pan.kang/ch368.git ${L4T_MODULE_CH36X_TOP}
			fi

			if [[ ! -f ${L4T_MODULE_CH36X} ]] || \
				[[ -f ${L4T_MODULE_CH36X} && $(get_module_vermagic ${L4T_MODULE_CH36X}) != ${L4T_KERNEL_VERMAGIC} ]];
			then
				make clean \
					-C ${L4T_MODULE_CH36X_TOP}/driver \
					-j$(nproc) \
					KERNEL_SRC=${L4T_SRC_MAKE_TOP} \
					KERNEL_OUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN}

				make \
					-C ${L4T_MODULE_CH36X_TOP}/driver \
					-j$(nproc) \
					KERNEL_SRC=${L4T_SRC_MAKE_TOP} \
					KERNEL_OUT=${L4T_SRC_OUTPUT_TOP} \
					ARCH=arm64 \
					CROSS_COMPILE=${L4T_TOOLCHAIN}

				local BUILD_RET=$?
				if [[ ${BUILD_RET} -ne 0 ]]; then
					echo -e "[ERROR] L4T Module CH36X build error!\n"
					exit 1
				fi
			fi
			;;
		"install")
			if [[ ! -f ${L4T_MODULE_CH36X} ]]; then
				${FUNCNAME} build

				if [[ ! -f ${L4T_MODULE_CH36X} ]]; then
					echo -e "[ERROR] L4T Module CH36X install error!\n"
					echo -e "CH36X Module file does not exist: ${L4T_MODULE_CH36X}.\n"
					exit 1
				fi
			fi

			echo -e "> ${L4T_MODULE_CH36X_ROOTFS_DIR}/$(basename ${L4T_MODULE_CH36X})"
			sudo_cmd "mkdir -p ${L4T_MODULE_CH36X_ROOTFS_DIR} \
				&& cp ${L4T_MODULE_CH36X} ${L4T_MODULE_CH36X_ROOTFS_DIR}"
			;;
		"uninstall")
			;;
		*)
			echo -e "[ERROR] ${FUNCNAME} argument error!\n"
			exit 1
			;;
	esac
}
